#include "AST/UniASTNode.h"
#include <utility>

void UniAST::TreeNode::addChild(UniAST::TreeNode *node)
{
    // fill children vector
    this->children.push_back(node);
    // fill ctx - child map
    if (node->ctx)
    {
        this->childrenCtxMap[node->ctx] = node;
    }
    // fill name - children map
    if (this->childrenNameMap.contains(node->nodeName))
    {
        this->childrenNameMap[node->nodeName].push_back(node);
    }
    else
    {
        this->childrenNameMap.insert(
            std::make_pair(node->nodeName, std::vector<TreeNode *>{node}));
    }
}
void UniAST::TreeNode::deleteChild(UniAST::TreeNode *node)
{
    // erase from children vector
    auto it =
        std::find(this->children.begin(), this->children.end(), node);
    if (it != this->children.end())
    {
        this->children.erase(it);
    }
    // erase from ctx - child map
    if (this->childrenCtxMap.contains(node->ctx))
    {
        this->childrenCtxMap.erase(node->ctx);
    }

    // erase from name - children map
    if (this->childrenNameMap.contains(node->nodeName))
    {
        // erase child
        auto &sameNamedChildren = this->childrenNameMap[node->nodeName];
        auto it = std::find(sameNamedChildren.begin(),
                            sameNamedChildren.end(), node);
        if (it != this->children.end())
        {
            this->children.erase(it);
        }
        // if vector is empty, delete this pair
        if (sameNamedChildren.empty())
        {
            this->childrenNameMap.erase(node->nodeName);
        }
    }
}

UniAST::TreeNode *
UniAST::TreeNode::findChild(antlr4::ParserRuleContext *ctx)
{
    if (this->childrenCtxMap.contains(ctx))
    {
        return childrenCtxMap[ctx];
    }
    return nullptr;
}

UniAST::TreeNode *UniAST::TreeNode::findChild(std::string name)
{
    if (this->childrenNameMap.contains(name))
    {
        return this->childrenNameMap[name][0];
    }
    return nullptr;
}

std::vector<UniAST::TreeNode *>
UniAST::TreeNode::findChildren(std::string name)
{
    if (this->childrenNameMap.contains(name))
    {
        return this->childrenNameMap[name];
    }
    return std::vector<UniAST::TreeNode *>{};
}

UniAST::TreeNode::TreeNode(antlr4::ParserRuleContext *ctx,
                           std::string nodeName)
    : ctx(ctx), nodeName(nodeName)
{
}

UniAST::TreeNode::~TreeNode()
{
    for (auto child : children)
    {
        delete child;
    }
}

std::string UniAST::TreeNode::toStringTree()
{
    std::string res = this->nodeName;
    res.append("<");
    res.append(this->ctx ? this->ctx->getText() : "generated");
    res.append(">");
    if (!this->children.empty())
    {
        res.append("(");
        for (auto child : this->children)
        {
            res.append(child->toStringTree());
            res.append(" ");
        }
        res.append(")");
    }
    return res;
}

std::string UniAST::TreeNode::getNodeName()
{
    return this->nodeName;
}

antlr4::ParserRuleContext *UniAST::TreeNode::getCtx()
{
    return this->ctx;
}

std::vector<UniAST::TreeNode *> &UniAST::TreeNode::getChildren()
{
    return this->children;
}

bool UniAST::RuleNode::isTerminal()
{
    return false;
}

UniAST::RuleNode::RuleNode(antlr4::ParserRuleContext *ctx,
                           std::string nodeName)
    : TreeNode(ctx, nodeName)
{
}

bool UniAST::TerminalNode::isTerminal()
{
    return true;
}

UniAST::TerminalNode::TerminalNode(antlr4::ParserRuleContext *ctx,
                                   std::string nodeName)
    : TreeNode(ctx, nodeName)
{
}